import json
import os
import sys
# 将包含 Model 的目录添加到系统路径
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
from algorithm.MR import Algorithm_MR
from utils.datasetUtils import *
from utils.modelUtils import getAvailableDevice, getResNet18
from utils.randomUtils import all_combinations
from _dynamic_dataset.AccCombinationRecord import AccCombinationRecord
from rootPath import project_path
from server_client.Server import Server
import torch
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
from datetime import datetime
'''
实验目的：验证各方数据集在相同大小和相同分布的情况下的真实夏普利值的计算
'''
# 获取当前日期和时间
now = datetime.now()
# 超参数
batch_size = 128
learning_rate = 0.0001
num_epochs = 200
# 读取配置文件
# 基于梯度的夏普利值计算
with open(project_path + "/conf/client.json", 'r') as f:
    clientConf = json.load(f)
with open(project_path + "/conf/server.json", 'r') as f:
    serverConf = json.load(f)
# # 配置日志记录器
dirPath = project_path + "/experiment/mr/logs/ddss" + f"/{now.year}_{now.month:02d}_{now.day:02d}_{now.hour:02d}_{now.minute:02d}_{now.second:02d}"
# # 创建目录（如果不存在）
# 数据预处理
train_transform = transforms.Compose(
    [
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
# 测试集数据预处理
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
device = getAvailableDevice()
# 服务器的验证集
test_set = torchvision.datasets.CIFAR10(root=project_path + '/data', train=False, download=True,
                                        transform=test_transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2)
model = getResNet18()
# 生成客户端的组合
clientNum = 5
clientIds = []
for i in range(clientNum):
    clientIds.append(i)
combinations = all_combinations(clientIds)
acc_record = AccCombinationRecord()
time_record = AccCombinationRecord()
# 服务器和客户端
server = Server(serverConf, test_loader, device, model)
# 创建相同分布相同大小的客户端
clients = get_DDSS_clients(clientConf, model, device, train_transform)
# 算法开始
Algorithm_MR(num_epochs, clients, server, device, combinations, dirPath, acc_record, time_record)
